#!/usr/bin/env python
"""

    run_muscle.py

"""

################################################################################
#
#   run_muscle
#
#
#   Copyright (C) 2009 Leo Goodstadt
#
#   This program is free software; you can redistribute it and/or
#   modify it under the terms of the GNU General Public License
#   as published by the Free Software Foundation; version 2
#   of the License
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program; if not, write to the Free Software
#   Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
#################################################################################

import sys, os
from collections import defaultdict
import json
#import popen2
import subprocess


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Functions        


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
def run_muscle (prot_sequences, muscle_bin = None):
    try:
        if not muscle_bin:
            muscle_bin = "/net/cpp-group/cluster/bin/muscle"
        proc = subprocess.Popen(muscle_bin + " -stable -quiet",
                                stdin=subprocess.PIPE,
                                stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE,
                                shell=True
                                )    
        for i, p in enumerate(prot_sequences):
            proc.stdin.write(">seq%d\n%s\n" % (i, p))
        proc.stdin.flush()
        cnt_sequences = len(prot_sequences)
        
        lines,stderr_str = proc.communicate()
        if "ERROR" in stderr_str:
            raise Exception("Failed to run muscle %s [%s]" % (stderr_str, str(prot_sequences)))
        aligned_seq = []
        for l in lines.split("\n"):
            if not len(l):
                continue
            if l[0] == ">":
                aligned_seq.append([])
                continue
            aligned_seq[-1].append(l)
    
    
        return ["".join(seq) for seq in aligned_seq]
    except:
        print prot_sequences, muscle_bin
        raise

        
import tempfile

def _write_temp_fasta_file (sequences, acc_prefix = "", filename_suffix = ".muscle"):
    """
    returns file name of written sequence
    """
    # write dna file
    sequence_file = tempfile.NamedTemporaryFile(suffix = filename_suffix, delete=False)
    sequence_file_name = sequence_file.name

    for i, p in enumerate(sequences):
        sequence_file.write(">seq_%s_%05d\n%s\n" % (acc_prefix, i, p))
    sequence_file.close()
    return sequence_file_name

        
def call_muscle (prot_sequences1, prot_sequences2, muscle_bin):
    """
    makes actual call. 
    Abstracted so we can call this twice if the first throws
    """
    file1 = _write_temp_fasta_file (prot_sequences1, acc_prefix = "A", filename_suffix = ".muscle")
    file2 = _write_temp_fasta_file (prot_sequences2, acc_prefix = "B", filename_suffix = ".muscle")
    

    cmds = muscle_bin + " -profile -in1 %s -in2 %s -quiet" % (file1, file2)
    process = subprocess.Popen(cmds, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)    
    lines, stderr_str = process.communicate()

    if process.returncode == 0:
        os.unlink(file1)
        os.unlink(file2)
        return lines
    else:
        process = subprocess.Popen(cmds, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)    
        lines, stderr_str = process.communicate()
        if process.returncode == 0:
            os.unlink(file1)
            os.unlink(file2)
            return lines
            
    raise Exception("Failed to run [%s]\n:%s\n" % (cmds,stderr_str))

from collections import defaultdict
def run_muscle_add_to_alignment (prot_sequences1, prot_sequences2, muscle_bin = None):
    if not muscle_bin:
        muscle_bin = "/net/cpp-group/cluster/bin/muscle"
    
    lines = call_muscle (prot_sequences1, prot_sequences2, muscle_bin)
    
                
    aligned_seq = defaultdict(list)
    curr_acc = None
    for l in lines.split("\n"):
        l = l.rstrip()
        if not len(l):
            continue
        if l[0] == ">":
            curr_acc = l[1:]
            continue
        aligned_seq[curr_acc].append(l)
        
    # turn each multi-line sequence into a single line
    
    aligned_seq = ["".join(seq) for acc, seq in sorted(aligned_seq.iteritems())]
    
    # split back into the original 2 sequences
    aligned_seq1 = aligned_seq[:len(prot_sequences1)]
    aligned_seq2 = aligned_seq[len(prot_sequences1):]
    if len(aligned_seq2) != len(prot_sequences2):
        raise Exception("Muscle failed with the input[%s] [%s] producing [%s]"
                            % (prot_sequences1, prot_sequences2, aligned_seq))
    return aligned_seq1, aligned_seq2
    
        
#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Testing


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
import unittest
class Test_run_muscle(unittest.TestCase):

    #       self.assertEqual(self.seq, range(10))
    #       self.assert_(element in self.seq)
    #       self.assertRaises(ValueError, random.sample, self.seq, 20)



    def test_function(self):
        """
            test 
        """
        prot_seq = [
        "HECARCGKNFSWHSDLILHEQIHSGEKP", 
        "HVCNECGKAFKTRNQLSMHRIIHTGEKP", 
        "FNCTQCGKAFNSRSALCRHKKTHSGEKP", 
        "HECRDCGKAFKTRNRLCMHQLIHTGEKP", 
        "YKCNCCGKAFQFKHSLTIHGRIHTGEKP", 
        "YECEECGKAFSGSSDLTKHIRIHTGERP", 
        "YECSKCGRAFSRSSDLSKHKRIHTREKH", 
        "YGCPQCGKDFSIKAELTKHRRIHTEEKR", 
        "YRCEECGKAFRHNCKRRAHEREHTGEKP", 
        "YQCRDCGKTFQDKHCLTIHQRIHTGEKP", 
        "YKCLECGKAFSGKSNLTNHRRIHTGEKP", 
        "HKCEVCGMAFHHSSVLRQHKRIHTGEKP", 
        "YTCSECGTSFRQGSALIGHKRVHTGEKP", 
        "YECEECGKAFRVSSNLTGHKKRKHQVWST"]
        print json.dumps(run_muscle(prot_seq), indent = 4)
        prot_seq = [
        'EHCSKCWGRYRVIRLSHVRLTVLIPILQMRELRHREVX', 
        'MTCPRSHSLVAELGFEPMTSDSQARTLS']
        print json.dumps(run_muscle(prot_seq), indent = 4)
        prot_seq = ['LRCNDGSSHRSTCPPPRSSPPX', 'ANCEMAHFKVLVSVLGSGETRRDDDIRG']
        print json.dumps(run_muscle(prot_seq), indent = 4)
        prot_seq = [     "HVCKECGKTFSHTSKLRKHQKFHTRLI",                                                                                                   
                         "HECSDCGKTFNHKLTLVHHQRIHTGKRL",                                                                                                  
                         "YECSECGKAFNNRSHLTRHEKVHTGERP",                                                                                                  
                         "FECSKCGRAFSQSSNFLRHQKVHTQVRP",                                                                                                  
                         "YGCNQCGKAFSRSSALIQHWRVHTGERP",                                                                                                  
                         "YECSECGRAFNNNSNLAQHQKVHTGERP",                                                                                                  
                         "FECSECGRNFSQSSHLLRHQKVHTGERP",                                                                                                  
                         "FECTECGKAFSNSSTLVQHQKVHTGKRP",                                                                                                  
                         "YECSECRKAFSRSSSLTQHWRIHTGERP",                                                                                                  
                         "YECSECGKAFAHSSSLIEHWRVHTRERP",                                                                                                  
                         "YECSECGKFFSQNSILTKHQKVHTGERP",                                                                                                  
                         "YECSECGKFFSRKSSLIYHWRVHTGERP",                                                                                                  
                         "YECSECGRGFSNNSHLVRHQRVHTQERP",                                                                                                  
                         "YECSQCGKFSESTLVRHQIVHTRERT"   ]
        first_seq = run_muscle_add_to_alignment([prot_seq[0]], prot_seq[1:-1])[0]
        last_seq  = run_muscle_add_to_alignment([prot_seq[-1]], prot_seq[1:-1])[0]
        print json.dumps((first_seq, prot_seq[1:-1], last_seq), indent = 4)


#
#   debug code not run if called as a module
#     
if __name__ == '__main__':
    if sys.argv.count("--debug"):
        sys.argv.remove("--debug")
    unittest.main()



    
